#! -*- encoding:utf-8 -*-
"""
@File    :   NDCG_metric.py
@Author  :   Zachary Li
@Contact :   li_zaaachary@163.com
@Dscpt   :   https://github.com/zygmuntz/evaluating-recommenders/blob/master/metrics.py

judgeScore 2: 准确或与问题高度相关; 1: 与问题相关; 0: 与问题无关

input type:
    data.json, data.jsonl 

example:
{
    "origin_question": xxxx,
    "candidates": [
        {
            "question": xxxxxxx,
            "judgeScore": x, (0, 1) 
        },
        {
            "question": xxxxxxx,
            "judgeScore": x, (0, 1) 
        },
        ...
    ]
}

"""
import argparse
import numpy as np
import json

# change this if using K > 100
denominator_table = np.log2( np.arange( 2, 102 ))


def dcg_at_k( r, k, method = 1 ):

	r = np.asfarray(r)[:k]
	if r.size:
		if method == 0:
			return r[0] + np.sum(r[1:] / np.log2(np.arange(2, r.size + 1)))
		elif method == 1:
			# return np.sum(r / np.log2(np.arange(2, r.size + 2)))
			return np.sum(r / denominator_table[:r.shape[0]])
		else:
			raise ValueError('method must be 0 or 1.')
	return 0.
 
 
def ndcg_at_k(r, k, method = 1 ):
    """
    2: 准确无误; 1: 与问题相关; 0: 与问题无关
    >>> r = [2, 2, 1, 0]
    """
    # import pdb; pdb.set_trace()
    idcg = dcg_at_k(sorted(r, reverse=True), k, method)
    # dcg_max = dcg_at_k(sorted(r, reverse=True), k, method)
    # dcg_min = dcg_at_k(sorted(r), k, method)
	# assert( dcg_max >= dcg_min )
    dcg_max = idcg
    if not dcg_max:
        return 0.
    
    dcg = dcg_at_k(r, k, method)
	
	#print dcg_min, dcg, dcg_max
	
    # 两种方法
    return dcg / idcg
    # return (dcg - dcg_min) / (dcg_max - dcg_min)


def load_data(file):
    with open(file, 'r', encoding='utf-8') as f:
        if "jsonl" in file:
            data = [json.loads(line) for line in f.readlines()]
        else:
            data = [json.load(f)]

    for example in data:
        # 根据系统排名排序
        # example['candidates'].sort(key=lambda x:x['whooshRank'])
        example['score_rank'] = [item['judgeScore'] for item in example['candidates']]
    
    # import pdb; pdb.set_trace()
    return data

def count_avg_ndcg(data):
    '''
    计算不同数量下的 NDCG
    '''
    method = 1
    total_avg = {1:0, 2:0, 3:0}
    for example in data:
        example['ndcg_score@1'] = ndcg_at_k(example['score_rank'], 1, method)
        total_avg[1] += example['ndcg_score@1']
        example['ndcg_score@2'] = ndcg_at_k(example['score_rank'], 2, method)
        total_avg[2] += example['ndcg_score@2']
        example['ndcg_score@3'] = ndcg_at_k(example['score_rank'], 3, method)
        total_avg[3] += example['ndcg_score@3']


    total_avg = {key: score/len(data) for key, score in total_avg.items()}
    return total_avg

def save_data(output_file, data, total_avg):
    f = open(output_file, 'w', encoding='utf-8')
    f.write(str(total_avg)+'\n')
    for example in data:
        f.write(json.dumps(example, ensure_ascii=False))
        f.write('\n')
    f.close()


if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description='根据数据文件计算NDCG指标')
    parser.add_argument('--input_file', type=str, help="请输入一个json(单一样本)或jsonl(每行一个样本)的路径")
    parser.add_argument('--output_file', type=str, default="result.jsonl", help="请输入一个json(单一样本)或jsonl(每行一个样本)的路径")

    # file = "data/data01.json"  "data/data02.jsonl"
    arg_str = r'''
    --input_file data\data03.json
    --output_file data\result.jsonl
    '''

    # args = parser.parse_args()
    args = parser.parse_args(arg_str.split())

    data = load_data(args.input_file)
    total_avg = count_avg_ndcg(data)
    print(total_avg)
    save_data(args.output_file, data, total_avg)